#include <CL/cl.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
#include <cassert>

#define CHECK_CL(err, msg) \
    if (err != CL_SUCCESS) { \
        std::cerr << "OpenCL error " << err << " at " << msg << std::endl; \
        exit(1); \
    }

std::string loadKernelSource(const std::string& filename) {
    std::string kernel_path = KERNELS_DIR "/" + filename;
    std::ifstream file(kernel_path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open kernel file: " + kernel_path);
    }
    return std::string((std::istreambuf_iterator<char>(file)), 
                      std::istreambuf_iterator<char>());
}

int main() {
    cl_int err;

    // 1️⃣ 获取平台和设备
    cl_uint numPlatforms = 0;
    CHECK_CL(clGetPlatformIDs(0, nullptr, &numPlatforms), "clGetPlatformIDs count");
    std::vector<cl_platform_id> platforms(numPlatforms);
    CHECK_CL(clGetPlatformIDs(numPlatforms, platforms.data(), nullptr), "clGetPlatformIDs list");

    cl_device_id device;
    CHECK_CL(clGetDeviceIDs(platforms[0], CL_DEVICE_TYPE_GPU, 1, &device, nullptr), "clGetDeviceIDs");

    // 2️⃣ 创建上下文与命令队列
    cl_context context = clCreateContext(nullptr, 1, &device, nullptr, nullptr, &err);
    CHECK_CL(err, "clCreateContext");
    cl_command_queue queue = clCreateCommandQueue(context, device, 0, &err);
    CHECK_CL(err, "clCreateCommandQueue");

    // 3️⃣ 读取并编译 kernel
    std::string kernelSrc = loadKernelSource("softmax.cl");
    const char *src = kernelSrc.c_str();
    size_t srcSize = kernelSrc.size();
    cl_program program = clCreateProgramWithSource(context, 1, &src, &srcSize, &err);
    CHECK_CL(err, "clCreateProgramWithSource");

    err = clBuildProgram(program, 1, &device, nullptr, nullptr, nullptr);
    if (err != CL_SUCCESS) {
        // 打印编译错误信息
        size_t logSize;
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr, &logSize);
        std::vector<char> log(logSize);
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, logSize, log.data(), nullptr);
        std::cerr << "Build failed:\n" << log.data() << std::endl;
        exit(1);
    }

    cl_kernel kernel = clCreateKernel(program, "kernel_soft_max", &err);
    CHECK_CL(err, "clCreateKernel");

    // 4️⃣ 构造测试数据
    const int ne00 = 4;
    std::vector<float> src0 = {1.0f, 2.0f, 3.0f, 4.0f};
    std::vector<float> dst(4, 0.0f);

    // 为简化测试，mask, src2 均为空
    cl_mem bufSrc0 = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                    sizeof(float) * src0.size(), src0.data(), &err);
    CHECK_CL(err, "clCreateBuffer src0");

    cl_mem bufDst = clCreateBuffer(context, CL_MEM_WRITE_ONLY,
                                   sizeof(float) * dst.size(), nullptr, &err);
    CHECK_CL(err, "clCreateBuffer dst");

    // 5️⃣ 设置 kernel 参数
    int arg = 0;
    cl_ulong zero = 0;
    cl_ulong nb = sizeof(float) * ne00;
    cl_int ne12 = 1, ne13 = 1, n_head_log2 = 0;
    float scale = 1.0f, max_bias = 0.0f, m0 = 0.0f, m1 = 0.0f;

    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_mem), &bufSrc0), "src0");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "offset0");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_mem), &bufSrc0), "src1");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "offset1");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_mem), &bufSrc0), "src2");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "offset2");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_mem), &bufDst),  "dst");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "offsetd");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_int), &ne00), "ne00");

    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &nb), "nb01");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "nb02");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "nb03");

    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_int), &ne12), "ne12");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_int), &ne13), "ne13");

    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "nb11");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "nb12");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "nb13");

    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &nb), "nb1");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "nb2");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_ulong), &zero), "nb3");

    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(float), &scale), "scale");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(float), &max_bias), "max_bias");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(float), &m0), "m0");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(float), &m1), "m1");
    CHECK_CL(clSetKernelArg(kernel, arg++, sizeof(cl_int), &n_head_log2), "n_head_log2");

    // 6️⃣ 执行 kernel
    size_t global = 4;
    size_t local = 4;
    err = clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, &global, &local, 0, nullptr, nullptr);
    CHECK_CL(err, "clEnqueueNDRangeKernel");

    // 7️⃣ 读取结果
    CHECK_CL(clEnqueueReadBuffer(queue, bufDst, CL_TRUE, 0, sizeof(float)*dst.size(), dst.data(), 0, nullptr, nullptr), "read dst");

    // 8️⃣ 打印结果
    std::cout << "Softmax result: ";
    for (auto v : dst) std::cout << v << " ";
    std::cout << std::endl;

    // 验证是否 sum≈1
    float sum = 0;
    for (auto v : dst) sum += v;
    std::cout << "Sum = " << sum << std::endl;

    // 9️⃣ 清理
    clReleaseMemObject(bufSrc0);
    clReleaseMemObject(bufDst);
    clReleaseKernel(kernel);
    clReleaseProgram(program);
    clReleaseCommandQueue(queue);
    clReleaseContext(context);
    return 0;
}
